import torch
import os
import argparse
import cv2
import torchvision
import numpy as np
import model as Model
from PIL import Image
import core.logger as Logger
import core.metrics as Metrics
import matplotlib.pyplot as plt
from torchvision.transforms import functional as trans_fn
from transformers import CLIPTextModel, CLIPTokenizer


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument('-c', '--config', type=str, default='config/sr_sr3_16_128.json',
                        help='JSON file for configuration')
    parser.add_argument('-p', '--phase', type=str, choices=['train', 'val'],
                        help='Run either train(training) or val(generation)', default='val')
    parser.add_argument('-gpu', '--gpu_ids', type=str, default=None)
    parser.add_argument('-debug', '-d', action='store_true')
    parser.add_argument('-enable_wandb', action='store_true')
    parser.add_argument('--img', type=str, default="dataset/celeba_test_16_128/lr_16/00052.png")
    parser.add_argument('--text', type=str, default="This person has eyeglasses, big nose, and straight hair.")
    parser.add_argument('--output_path', type=str, default="./output")
    parser.add_argument('--sizes', type=str, default='16, 128')

    args = parser.parse_args()
    opt = Logger.parse(args)
    opt = Logger.dict_to_nonedict(opt)

    # images and prompts
    img = Image.open(args.img)
    img = img.convert('RGB')
    size = opt['datasets']['val']['r_resolution']

    def resize_and_convert(img, size, resample):
        if (img.size[0] != size):
            img = trans_fn.resize(img, size, resample)
            img = trans_fn.center_crop(img, size)
        return img

    totensor = torchvision.transforms.ToTensor()
    img = totensor(resize_and_convert(img, size, Image.Resampling.BICUBIC)) * 2 - 1
    img = img.unsqueeze(0).to("cuda")

    prompt = args.text

    # model
    text_encoder = CLIPTextModel.from_pretrained("pretrained/text_encoder").to("cuda")
    tokenizer = CLIPTokenizer.from_pretrained("pretrained/tokenizer")
    diffusion = Model.create_model(opt)

    # infer
    with torch.no_grad():
        text_input_ids = tokenizer(
            prompt,
            padding="max_length",
            max_length=tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt"
        ).input_ids
        text_embeddings = torch.clamp(text_encoder(text_input_ids.to("cuda"))[0], -1.0, 1.0)

    data = {"SR": img, 'Prompt': text_embeddings}
    diffusion.feed_data(data)
    diffusion.test(continous=False)
    visuals = diffusion.get_current_visuals(need_LR=False, sample=True)


    result_path = '{}'.format(opt['path']['results'])

    sr_img = visuals['SAM']
    Metrics.save_img(Metrics.tensor2img(sr_img), '{}/sr_.png'.format(result_path))

#     else:
#         # grid img
#         sr_img = Metrics.tensor2img(visuals['SR'])  # uint8
#         Metrics.save_img(
#             sr_img, '{}/{}_{}_sr_process.png'.format(result_path, current_step, idx))
#         Metrics.save_img(
#             Metrics.tensor2img(visuals['SR'][-1]), '{}/{}_{}_sr.png'.format(result_path, current_step, idx))
#
#     Metrics.save_img(
#         hr_img, '{}/{}_{}_hr.png'.format(result_path, current_step, idx))
#     Metrics.save_img(
#         fake_img, '{}/{}_{}_inf.png'.format(result_path, current_step, idx))
#
#     if wandb_logger and opt['log_infer']:
#         wandb_logger.log_eval_data(fake_img, Metrics.tensor2img(visuals['SR'][-1]), hr_img)
#
# if wandb_logger and opt['log_infer']:
#     wandb_logger.log_eval_table(commit=True)


